
from openai import OpenAI
import json
from src.utils import get_project_root


class ElementAgent:

    def __init__(self, new_dom_data: str = None, old_dom_data: str = None, page_name: str = None,
                 new_dom_data_path=None, old_dom_data_path=None):
        # 传文件地址
        self.client = OpenAI(
            # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
            api_key='sk-8dcd9ec46bbe42f3947c71d27a21774d',
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        self.new_dom_data = new_dom_data
        self.old_dom_data = old_dom_data
        self.new_dom_data_path = new_dom_data_path
        self.old_dom_data_path = old_dom_data_path
        self.page_name = page_name
        self.tools = [
            {
                "type": "function",
                "function": {
                    "name": "write_to_file",
                    "description": "将给定的内容写入指定的文件。",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "filename": {
                                "type": "string",
                                "description": "要写入的文件名",
                            },
                            "content": {
                                "type": "string",
                                "description": "要写入文件的内容"
                            }
                        },
                        "required": ["filename", "content"],
                    },
                }
            }
        ]
        system_prompt_path = get_project_root() / 'config' / "login_page_general_system_prompt.md"
        with open(system_prompt_path, "r", encoding="utf-8") as f:
            self.system_prompt = f.read()

        if new_dom_data_path and new_dom_data is None:
            with open(self.new_dom_data_path, "r", encoding="utf-8") as f:
                self.new_dom_data = f.read()
        elif new_dom_data and new_dom_data_path is None:
            self.new_dom_data = new_dom_data
        else:
            raise "new_dom_data_path or new_dom_data must be provided"
        if old_dom_data_path and old_dom_data is None:
            with open(self.old_dom_data_path, "r", encoding="utf-8") as f:
                self.old_dom_data_path = f.read()
        elif old_dom_data and old_dom_data_path is None:
            self.old_dom_data = old_dom_data
        elif old_dom_data is None and old_dom_data_path is None:
            self.old_dom_data = None


    def add_tools(self):
        # Step 1: 定义 tools（函数描述）
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "write_to_file",
                    "description": "将给定的内容写入指定的文件。",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "filename": {
                                "type": "string",
                                "description": "要写入的文件名",
                            },
                            "content": {
                                "type": "string",
                                "description": "要写入文件的内容"
                            }
                        },
                        "required": ["filename", "content"],
                    },
                }
            }
        ]
        self.tools.append(tools[0])

    # Step 2: 实现工具函数

    def call_model(self):
        completion = self.client.chat.completions.create(
            model="qwen3-coder-plus-2025-07-22",
            # Model list: https://help.aliyun.com/zh/model-studio/getting-started/models
            messages=[
                {'role': 'system', 'content': "You are a helpful assistant."},  # Fixed missing closing brace
                {'role': 'user',
                 'content': f'{self.system_prompt},page_name:{self.page_name},new_dom_data:{self.new_dom_data},old_dom_data:{self.old_dom_data}'}
            ],
            tool_choice={"type": "function", "function": {"name": "write_to_file"}},
            tools=self.tools,
        )
        response = completion.model_dump()
        print("=" * 50)
        print(response)
        model_response_content = None

        if 'choices' in response and len(response['choices']) > 0:
            first_choice = response['choices'][0]
            if 'message' in first_choice and 'content' in first_choice['message']:
                model_response_content = first_choice['message']

        # 打印模型生成的内容
        if model_response_content is not None:
            print("=" * 50)
            print("模型返回的结果：")
            print(model_response_content)
        else:
            print("未找到模型返回的内容。")

        if model_response_content.get('tool_calls'):
            for tool_call in model_response_content['tool_calls']:
                function_name = tool_call['function']['name']
                arguments = json.loads(tool_call['function']['arguments'])

                if function_name == "write_to_file":
                    # 调用工具函数

                    result = write_to_file(arguments['filename'], arguments['content'])

                    # print(f"Tool call result: {result}")
                    self.client.close()
                    return result


def write_to_file(filename: str, content: str):
    """将给定的内容写入指定的文件"""
    proj_root = get_project_root()
    file_path = proj_root / 'config' / filename
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)
        return {"status": "success", "message": f"Content written to {filename}"}
    except Exception as e:
        return {"status": "error", "message": str(e)}


if __name__ == '__main__':
    old_file = r'/src/config\page.html'
    new_file = r'D:\Programming_soft\browser-use\playwright_demo\pages\login_page\page.html'
    print(ElementAgent(old_dom_data=old_file, new_dom_data=new_file).call_model())
